package com.ae.shield.generator;

import com.ae.shield.common.block.BlockList;
import com.ae.shield.common.items.ItemList;
import com.google.common.collect.ImmutableList;
import com.mojang.datafixers.util.Pair;
import net.minecraft.block.Block;
import net.minecraft.block.FlowingFluidBlock;
import net.minecraft.data.DataGenerator;
import net.minecraft.data.loot.BlockLootTables;
import net.minecraft.data.loot.EntityLootTables;
import net.minecraft.data.loot.FishingLootTables;
import net.minecraft.entity.EntityType;
import net.minecraft.loot.*;
import net.minecraft.loot.conditions.KilledByPlayer;
import net.minecraft.loot.conditions.RandomChanceWithLooting;
import net.minecraft.util.ResourceLocation;
import net.minecraftforge.common.data.ForgeLootTableProvider;
import net.minecraftforge.fml.RegistryObject;
import net.minecraftforge.registries.ForgeRegistries;

import java.util.ArrayList;
import java.util.List;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Supplier;

public class ModLootTableProvider extends ForgeLootTableProvider {
    public ModLootTableProvider(DataGenerator gen) {
        super(gen);
    }

    @Override
    protected List<Pair<Supplier<Consumer<BiConsumer<ResourceLocation, LootTable.Builder>>>, LootParameterSet>> getTables() {
        return ImmutableList.of(
                Pair.of(ModBlockLootTable::new, LootParameterSets.BLOCK),
                Pair.of(ModEntityLootTable::new, LootParameterSets.ENTITY),
                Pair.of(ModFishLootTable::new, LootParameterSets.FISHING));
    }

    private static class ModBlockLootTable extends BlockLootTables {
        @Override
        protected void addTables() {
            for (RegistryObject<Block> entry : BlockList.BLOCKS.getEntries()) {
                if (entry.get() instanceof FlowingFluidBlock) {
                    continue;
                }

                registerDropSelfLootTable(entry.get());
            }
        }

        @Override
        protected Iterable<Block> getKnownBlocks() {
            List<Block> blockList = new ArrayList<>();
            for (RegistryObject<Block> registry : BlockList.BLOCKS.getEntries()) {

                if (registry.get() instanceof FlowingFluidBlock) {
                    continue;
                }

                if (ForgeRegistries.ITEMS.containsKey(registry.get().getRegistryName())) {
                    blockList.add(registry.get());
                }
            }
            return blockList;
        }
    }

    private static class ModEntityLootTable extends EntityLootTables {
        @Override
        protected void addTables() {
            registerLootTable(EntityType.VINDICATOR, LootTable.builder()
                    .addLootPool(LootPool.builder()
                            .rolls(ConstantRange.of(1))
                            .addEntry(ItemLootEntry.builder(ItemList.TRANSFORMED_TOTEM.get()))
                            .acceptCondition(KilledByPlayer.builder())
                            .acceptCondition(RandomChanceWithLooting.builder(0.01f, 0.01f))));
            registerLootTable(EntityType.EVOKER, LootTable.builder()
                    .addLootPool(LootPool.builder()
                            .rolls(ConstantRange.of(1))
                            .addEntry(ItemLootEntry.builder(ItemList.TRANSFORMED_TOTEM.get()))
                            .acceptCondition(KilledByPlayer.builder())
                            .acceptCondition(RandomChanceWithLooting.builder(0.015f, 0.01f))));
        }

        @Override
        protected Iterable<EntityType<?>> getKnownEntities() {
            ArrayList<EntityType<?>> list = new ArrayList<>();
            list.add(EntityType.VINDICATOR);
            list.add(EntityType.EVOKER);
            return list;
        }
    }

    private static class ModFishLootTable extends FishingLootTables {
        @Override
        public void accept(BiConsumer<ResourceLocation, LootTable.Builder> table) {
            table.accept(LootTables.GAMEPLAY_FISHING_TREASURE, LootTable.builder()
                    .addLootPool(LootPool.builder()
                            .rolls(ConstantRange.of(1))
                            .addEntry(ItemLootEntry.builder(ItemList.TRANSFORMED_TOTEM.get()))
                            .acceptCondition(RandomChanceWithLooting.builder(0.001f, 0.001f))));
        }
    }
}
